package cn.maihe.elg.operation.utils;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Arrays;
import java.util.Base64;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;

public class AESFileUtil {

    /**
     * init AES Cipher
     *
     * @param passsword
     * @param cipherMode
     * @return
     */
    public static Cipher initAESCipher(String passsword, int cipherMode) {
        Cipher cipher = null;
        try {
            SecretKey key = getAesSecretKey(passsword);
            cipher = Cipher.getInstance("AES/ECB/PKCS7Padding");
            cipher.init(cipherMode, key);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return cipher;
    }

    private static SecretKey getAesSecretKey(String password) {
        KeyGenerator keyGenerator;
        SecureRandom secureRandom = null;
        try {
            keyGenerator = KeyGenerator.getInstance("AES");
            secureRandom = SecureRandom.getInstance("SHA1PRNG");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        secureRandom.setSeed(password.getBytes());
        keyGenerator.init(128, secureRandom);
        SecretKey secretKey = keyGenerator.generateKey();
        return secretKey;
    }

    private static SecretKey getSecretKey(String password) {
        int keyLength = 256;
        byte[] keyBytes = new byte[keyLength / 8];
        SecretKeySpec key = null;
        try {
            Arrays.fill(keyBytes, (byte) 0x0);
            Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
            byte[] passwordBytes = password.getBytes("UTF-8");
            int length = passwordBytes.length < keyBytes.length ? passwordBytes.length : keyBytes.length;
            System.arraycopy(passwordBytes, 0, keyBytes, 0, length);

            key = new SecretKeySpec(keyBytes, "AES");
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }
        return key;
    }


    /**
     * AES 加密 (返回字节数组)
     *
     * @param encryptPath 要加密的文件路径
     * @param sKey        加密密钥
     * @return 加密后的字节数组
     */
    public static byte[] encryptFileToByteArray(String encryptPath, String sKey) {
        File encryptFile;
        try {
            encryptFile = new File(encryptPath);
            if (!encryptFile.exists()) {
                throw new NullPointerException("Encrypt file is empty");
            }
            Cipher cipher = initAESCipher(sKey, Cipher.ENCRYPT_MODE);

            try (FileInputStream fileInputStream = new FileInputStream(encryptFile);
                 ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream()) {

                byte[] buffer = new byte[1024];
                int bufferLength;

                while ((bufferLength = fileInputStream.read(buffer)) != -1) {
                    byte[] encryptedChunk = cipher.update(buffer, 0, bufferLength);
                    if (encryptedChunk != null) {
                        byteArrayOutputStream.write(encryptedChunk);
                    }
                }
                byte[] finalEncryptedChunk = cipher.doFinal();
                if (finalEncryptedChunk != null) {
                    byteArrayOutputStream.write(finalEncryptedChunk);
                }

                return byteArrayOutputStream.toByteArray();
            }
        } catch (IOException | IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }


    /**
     * AES 加密
     *
     * @param encryptPath
     * @param decryptPath
     * @param sKey
     * @return
     */
    public static boolean encryptFile(String encryptPath, String decryptPath, String sKey) {
        File encryptFile = null;
        File decryptfile = null;
        CipherOutputStream cipherOutputStream = null;
        BufferedInputStream bufferedInputStream = null;
        try {
            encryptFile = new File(encryptPath);
            if (!encryptFile.exists()) {
                throw new NullPointerException("Encrypt file is empty");
            }
            decryptfile = new File(decryptPath);
            if (decryptfile.exists()) {
                decryptfile.delete();
            }
            decryptfile.createNewFile();

            Cipher cipher = initAESCipher(sKey, Cipher.ENCRYPT_MODE);
            cipherOutputStream = new CipherOutputStream(new FileOutputStream(decryptfile), cipher);
            bufferedInputStream = new BufferedInputStream(new FileInputStream(encryptFile));

            byte[] buffer = new byte[1024];
            int bufferLength;

            while ((bufferLength = bufferedInputStream.read(buffer)) != -1) {
                cipherOutputStream.write(buffer, 0, bufferLength);
            }

            System.out.println("Encrypt file 1, size: " + buffer.length + "   " + Arrays.toString(buffer));

            bufferedInputStream.close();
            cipherOutputStream.close();
//            delFile(encryptPath);
        } catch (IOException e) {
            e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
            return false;
        }
        return true;
    }

    /**
     * AES 解密
     *
     * @param encryptPath
     * @param decryptPath
     * @param mKey
     * @return
     */
    public static boolean decryptFile(String encryptPath, String decryptPath, String mKey) {
        File encryptFile;
        File decryptFile = null;
        BufferedOutputStream outputStream = null;
        CipherInputStream inputStream = null;
        try {
            encryptFile = new File(encryptPath);
            if (!encryptFile.exists()) {
                throw new NullPointerException("Decrypt file is empty");
            }
            decryptFile = new File(decryptPath);
            if (decryptFile.exists()) {
                decryptFile.delete();
            }
            decryptFile.createNewFile();

            Cipher cipher = initAESCipher(mKey, Cipher.DECRYPT_MODE);

            outputStream = new BufferedOutputStream(Files.newOutputStream(decryptFile.toPath()));
            inputStream = new CipherInputStream(Files.newInputStream(encryptFile.toPath()), cipher);

            int bufferLength;
            byte[] buffer = new byte[1024];

            while ((bufferLength = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bufferLength);
            }
            inputStream.close();
            outputStream.close();
//            delFile(encryptPath);
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }


    /**
     * AES 加密字符串
     *
     * @param plainText 要加密的字符串
     * @param sKey      加密密钥
     * @return 加密后的 Base64 编码字符串
     */
    public static String encryptString(String plainText, String sKey) {
        try {
            Cipher cipher = initAESCipher(sKey, Cipher.ENCRYPT_MODE);
            byte[] encryptedBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(encryptedBytes);
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * AES 解密字符串
     *
     * @param encryptedText 要解密的 Base64 编码字符串
     * @param sKey          解密密钥
     * @return 解密后的字符串
     */
    public static String decryptString(String encryptedText, String sKey) {
        try {
            Cipher cipher = initAESCipher(sKey, Cipher.DECRYPT_MODE);
            byte[] decodedBytes = Base64.getDecoder().decode(encryptedText);
            byte[] decryptedBytes = cipher.doFinal(decodedBytes);
            return new String(decryptedBytes, StandardCharsets.UTF_8);
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }


    /**
     * AES 加密字符串（返回字节数组）
     *
     * @param plainText 要加密的字符串
     * @param sKey      加密密钥
     * @return 加密后的字节数组
     */
    public static byte[] encryptStringToByteArray(String plainText, String sKey) {
        try {
            Cipher cipher = initAESCipher(sKey, Cipher.ENCRYPT_MODE);
            return cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }
}
